{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/home/soda/rcappuzz/work/benchmark-join-suggestions\n"
     ]
    }
   ],
   "source": [
    "cd ~/bench"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import polars as pl\n",
    "from pathlib import Path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "raw_data_path = Path(\"data/source_tables/raw\")\n",
    "yadl_data_path = Path(\"data/source_tables/yadl\")\n",
    "od_data_path = Path(\"data/source_tables/open_data_us\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Movies improved"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "import ast\n",
    "\n",
    "\n",
    "def clean_genres(ll):\n",
    "    g = ast.literal_eval(ll)\n",
    "    try:\n",
    "        l1 = g[0][\"name\"]\n",
    "        return l1\n",
    "    except IndexError:\n",
    "        return \"\"\n",
    "\n",
    "\n",
    "def clean_production_companies(ll):\n",
    "    try:\n",
    "        g = ast.literal_eval(ll)\n",
    "    except ValueError:\n",
    "        return \"\"\n",
    "    except SyntaxError:\n",
    "        print(ll)\n",
    "    try:\n",
    "        l1 = g[0][\"name\"]\n",
    "        return l1\n",
    "    except IndexError:\n",
    "        return \"\"\n",
    "    except TypeError:\n",
    "        return \"\"\n",
    "\n",
    "\n",
    "def clean_production_country(ll):\n",
    "    try:\n",
    "        g = ast.literal_eval(ll)\n",
    "    except ValueError:\n",
    "        return \"\"\n",
    "    try:\n",
    "        l1 = g[0][\"iso_3166_1\"]\n",
    "        return l1\n",
    "    except IndexError:\n",
    "        return \"\"\n",
    "    except TypeError:\n",
    "        return \"\"\n",
    "\n",
    "\n",
    "def clean_spoken_language(ll):\n",
    "    try:\n",
    "        g = ast.literal_eval(ll)\n",
    "    except ValueError:\n",
    "        return \"\"\n",
    "    try:\n",
    "        l1 = g[0][\"name\"]\n",
    "        return l1\n",
    "    except IndexError:\n",
    "        return \"\"\n",
    "    except TypeError:\n",
    "        return \"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "df = pl.read_parquet(\"data/source_tables/movie_revenues.parquet\").to_pandas()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "df = df.drop(\n",
    "    [\n",
    "        \"belongs_to_collection\",\n",
    "        \"homepage\",\n",
    "        \"imdb_id\",\n",
    "        \"overview\",\n",
    "        \"tagline\",\n",
    "        \"poster_path\",\n",
    "        \"release_date\",\n",
    "    ],\n",
    "    axis=1,\n",
    ")\n",
    "\n",
    "df.genres = df.genres.apply(clean_genres)\n",
    "df.production_companies = df.production_companies.apply(clean_production_companies)\n",
    "df.production_countries = df.production_countries.apply(clean_production_country)\n",
    "df.spoken_languages = df.spoken_languages.apply(clean_spoken_language)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "df = df.drop([\"yago4_col_to_embed\", \"raw_entities\"], axis=1).rename(\n",
    "    {\"yago3_col_to_embed\": \"col_to_embed\"}, axis=1\n",
    ")\n",
    "\n",
    "df = pl.from_pandas(df).drop_nulls(\"col_to_embed\").to_pandas()\n",
    "df[\"col_to_embed\"] = \"<\" + df[\"col_to_embed\"] + \">\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(7397, 20)"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>adult</th>\n",
       "      <th>budget</th>\n",
       "      <th>genres</th>\n",
       "      <th>id</th>\n",
       "      <th>original_language</th>\n",
       "      <th>original_title</th>\n",
       "      <th>popularity</th>\n",
       "      <th>production_companies</th>\n",
       "      <th>production_countries</th>\n",
       "      <th>revenue</th>\n",
       "      <th>runtime</th>\n",
       "      <th>spoken_languages</th>\n",
       "      <th>status</th>\n",
       "      <th>title</th>\n",
       "      <th>video</th>\n",
       "      <th>vote_average</th>\n",
       "      <th>vote_count</th>\n",
       "      <th>year</th>\n",
       "      <th>col_to_embed</th>\n",
       "      <th>target</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>False</td>\n",
       "      <td>30000000</td>\n",
       "      <td>Animation</td>\n",
       "      <td>862</td>\n",
       "      <td>en</td>\n",
       "      <td>Toy Story</td>\n",
       "      <td>21.946943</td>\n",
       "      <td>Pixar Animation Studios</td>\n",
       "      <td>US</td>\n",
       "      <td>373554033.0</td>\n",
       "      <td>81.0</td>\n",
       "      <td>English</td>\n",
       "      <td>Released</td>\n",
       "      <td>Toy Story</td>\n",
       "      <td>False</td>\n",
       "      <td>7.7</td>\n",
       "      <td>5415.0</td>\n",
       "      <td>1995</td>\n",
       "      <td>&lt;Toy_Story&gt;</td>\n",
       "      <td>8.572353</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>False</td>\n",
       "      <td>65000000</td>\n",
       "      <td>Adventure</td>\n",
       "      <td>8844</td>\n",
       "      <td>en</td>\n",
       "      <td>Jumanji</td>\n",
       "      <td>17.015539</td>\n",
       "      <td>TriStar Pictures</td>\n",
       "      <td>US</td>\n",
       "      <td>262797249.0</td>\n",
       "      <td>104.0</td>\n",
       "      <td>English</td>\n",
       "      <td>Released</td>\n",
       "      <td>Jumanji</td>\n",
       "      <td>False</td>\n",
       "      <td>6.9</td>\n",
       "      <td>2413.0</td>\n",
       "      <td>1995</td>\n",
       "      <td>&lt;Jumanji&gt;</td>\n",
       "      <td>8.419621</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>False</td>\n",
       "      <td>16000000</td>\n",
       "      <td>Comedy</td>\n",
       "      <td>31357</td>\n",
       "      <td>en</td>\n",
       "      <td>Waiting to Exhale</td>\n",
       "      <td>3.859495</td>\n",
       "      <td>Twentieth Century Fox Film Corporation</td>\n",
       "      <td>US</td>\n",
       "      <td>81452156.0</td>\n",
       "      <td>127.0</td>\n",
       "      <td>English</td>\n",
       "      <td>Released</td>\n",
       "      <td>Waiting to Exhale</td>\n",
       "      <td>False</td>\n",
       "      <td>6.1</td>\n",
       "      <td>34.0</td>\n",
       "      <td>1995</td>\n",
       "      <td>&lt;Waiting_to_Exhale&gt;</td>\n",
       "      <td>7.910903</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>False</td>\n",
       "      <td>0</td>\n",
       "      <td>Comedy</td>\n",
       "      <td>11862</td>\n",
       "      <td>en</td>\n",
       "      <td>Father of the Bride Part II</td>\n",
       "      <td>8.387519</td>\n",
       "      <td>Sandollar Productions</td>\n",
       "      <td>US</td>\n",
       "      <td>76578911.0</td>\n",
       "      <td>106.0</td>\n",
       "      <td>English</td>\n",
       "      <td>Released</td>\n",
       "      <td>Father of the Bride Part II</td>\n",
       "      <td>False</td>\n",
       "      <td>5.7</td>\n",
       "      <td>173.0</td>\n",
       "      <td>1995</td>\n",
       "      <td>&lt;Father_of_the_Bride_Part_II&gt;</td>\n",
       "      <td>7.884109</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>False</td>\n",
       "      <td>60000000</td>\n",
       "      <td>Action</td>\n",
       "      <td>949</td>\n",
       "      <td>en</td>\n",
       "      <td>Heat</td>\n",
       "      <td>17.924927</td>\n",
       "      <td>Regency Enterprises</td>\n",
       "      <td>US</td>\n",
       "      <td>187436818.0</td>\n",
       "      <td>170.0</td>\n",
       "      <td>English</td>\n",
       "      <td>Released</td>\n",
       "      <td>Heat</td>\n",
       "      <td>False</td>\n",
       "      <td>7.7</td>\n",
       "      <td>1886.0</td>\n",
       "      <td>1995</td>\n",
       "      <td>&lt;Heat_(1995_film)&gt;</td>\n",
       "      <td>8.272855</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7392</th>\n",
       "      <td>False</td>\n",
       "      <td>750000</td>\n",
       "      <td>Crime</td>\n",
       "      <td>280422</td>\n",
       "      <td>ru</td>\n",
       "      <td>Все и сразу</td>\n",
       "      <td>0.201582</td>\n",
       "      <td>Кинокомпания «Lunapark»</td>\n",
       "      <td>RU</td>\n",
       "      <td>3.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>Pусский</td>\n",
       "      <td>Released</td>\n",
       "      <td>All at Once</td>\n",
       "      <td>False</td>\n",
       "      <td>6.0</td>\n",
       "      <td>4.0</td>\n",
       "      <td>2014</td>\n",
       "      <td>&lt;All_at_Once_(film)&gt;</td>\n",
       "      <td>0.477121</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7393</th>\n",
       "      <td>False</td>\n",
       "      <td>0</td>\n",
       "      <td>Drama</td>\n",
       "      <td>240789</td>\n",
       "      <td>ru</td>\n",
       "      <td>Чудо</td>\n",
       "      <td>0.436028</td>\n",
       "      <td>Central Partnership</td>\n",
       "      <td>RU</td>\n",
       "      <td>50656.0</td>\n",
       "      <td>110.0</td>\n",
       "      <td>Pусский</td>\n",
       "      <td>Released</td>\n",
       "      <td>The Miracle</td>\n",
       "      <td>False</td>\n",
       "      <td>6.3</td>\n",
       "      <td>3.0</td>\n",
       "      <td>2009</td>\n",
       "      <td>&lt;The_Miracle_(2009_film)&gt;</td>\n",
       "      <td>4.704631</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7394</th>\n",
       "      <td>False</td>\n",
       "      <td>800000</td>\n",
       "      <td>Comedy</td>\n",
       "      <td>62757</td>\n",
       "      <td>en</td>\n",
       "      <td>Dikari</td>\n",
       "      <td>0.903061</td>\n",
       "      <td></td>\n",
       "      <td></td>\n",
       "      <td>1328612.0</td>\n",
       "      <td>100.0</td>\n",
       "      <td>Pусский</td>\n",
       "      <td>Released</td>\n",
       "      <td>Savages</td>\n",
       "      <td>False</td>\n",
       "      <td>5.8</td>\n",
       "      <td>6.0</td>\n",
       "      <td>2006</td>\n",
       "      <td>&lt;Savages_(2006_film)&gt;</td>\n",
       "      <td>6.123398</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7395</th>\n",
       "      <td>False</td>\n",
       "      <td>2000000</td>\n",
       "      <td>Romance</td>\n",
       "      <td>63281</td>\n",
       "      <td>en</td>\n",
       "      <td>Про любоff</td>\n",
       "      <td>0.121844</td>\n",
       "      <td>Profit</td>\n",
       "      <td>RU</td>\n",
       "      <td>1268793.0</td>\n",
       "      <td>107.0</td>\n",
       "      <td>Pусский</td>\n",
       "      <td>Released</td>\n",
       "      <td>Pro Lyuboff</td>\n",
       "      <td>False</td>\n",
       "      <td>4.0</td>\n",
       "      <td>3.0</td>\n",
       "      <td>2010</td>\n",
       "      <td>&lt;Pro_Lyuboff_(2010_film)&gt;</td>\n",
       "      <td>6.103391</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7396</th>\n",
       "      <td>False</td>\n",
       "      <td>5000000</td>\n",
       "      <td>Action</td>\n",
       "      <td>63898</td>\n",
       "      <td>ru</td>\n",
       "      <td>Антидурь</td>\n",
       "      <td>0.039793</td>\n",
       "      <td></td>\n",
       "      <td>RU</td>\n",
       "      <td>1413000.0</td>\n",
       "      <td>91.0</td>\n",
       "      <td>Pусский</td>\n",
       "      <td>Released</td>\n",
       "      <td>Antidur</td>\n",
       "      <td>False</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>2007</td>\n",
       "      <td>&lt;Antidur_(2007_film)&gt;</td>\n",
       "      <td>6.150142</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>7397 rows × 20 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "      adult    budget     genres      id original_language  \\\n",
       "0     False  30000000  Animation     862                en   \n",
       "1     False  65000000  Adventure    8844                en   \n",
       "2     False  16000000     Comedy   31357                en   \n",
       "3     False         0     Comedy   11862                en   \n",
       "4     False  60000000     Action     949                en   \n",
       "...     ...       ...        ...     ...               ...   \n",
       "7392  False    750000      Crime  280422                ru   \n",
       "7393  False         0      Drama  240789                ru   \n",
       "7394  False    800000     Comedy   62757                en   \n",
       "7395  False   2000000    Romance   63281                en   \n",
       "7396  False   5000000     Action   63898                ru   \n",
       "\n",
       "                   original_title  popularity  \\\n",
       "0                       Toy Story   21.946943   \n",
       "1                         Jumanji   17.015539   \n",
       "2               Waiting to Exhale    3.859495   \n",
       "3     Father of the Bride Part II    8.387519   \n",
       "4                            Heat   17.924927   \n",
       "...                           ...         ...   \n",
       "7392                  Все и сразу    0.201582   \n",
       "7393                         Чудо    0.436028   \n",
       "7394                       Dikari    0.903061   \n",
       "7395                   Про любоff    0.121844   \n",
       "7396                     Антидурь    0.039793   \n",
       "\n",
       "                        production_companies production_countries  \\\n",
       "0                    Pixar Animation Studios                   US   \n",
       "1                           TriStar Pictures                   US   \n",
       "2     Twentieth Century Fox Film Corporation                   US   \n",
       "3                      Sandollar Productions                   US   \n",
       "4                        Regency Enterprises                   US   \n",
       "...                                      ...                  ...   \n",
       "7392                 Кинокомпания «Lunapark»                   RU   \n",
       "7393                     Central Partnership                   RU   \n",
       "7394                                                                \n",
       "7395                                  Profit                   RU   \n",
       "7396                                                           RU   \n",
       "\n",
       "          revenue  runtime spoken_languages    status  \\\n",
       "0     373554033.0     81.0          English  Released   \n",
       "1     262797249.0    104.0          English  Released   \n",
       "2      81452156.0    127.0          English  Released   \n",
       "3      76578911.0    106.0          English  Released   \n",
       "4     187436818.0    170.0          English  Released   \n",
       "...           ...      ...              ...       ...   \n",
       "7392          3.0      0.0          Pусский  Released   \n",
       "7393      50656.0    110.0          Pусский  Released   \n",
       "7394    1328612.0    100.0          Pусский  Released   \n",
       "7395    1268793.0    107.0          Pусский  Released   \n",
       "7396    1413000.0     91.0          Pусский  Released   \n",
       "\n",
       "                            title  video  vote_average  vote_count  year  \\\n",
       "0                       Toy Story  False           7.7      5415.0  1995   \n",
       "1                         Jumanji  False           6.9      2413.0  1995   \n",
       "2               Waiting to Exhale  False           6.1        34.0  1995   \n",
       "3     Father of the Bride Part II  False           5.7       173.0  1995   \n",
       "4                            Heat  False           7.7      1886.0  1995   \n",
       "...                           ...    ...           ...         ...   ...   \n",
       "7392                  All at Once  False           6.0         4.0  2014   \n",
       "7393                  The Miracle  False           6.3         3.0  2009   \n",
       "7394                      Savages  False           5.8         6.0  2006   \n",
       "7395                  Pro Lyuboff  False           4.0         3.0  2010   \n",
       "7396                      Antidur  False           1.0         1.0  2007   \n",
       "\n",
       "                       col_to_embed    target  \n",
       "0                       <Toy_Story>  8.572353  \n",
       "1                         <Jumanji>  8.419621  \n",
       "2               <Waiting_to_Exhale>  7.910903  \n",
       "3     <Father_of_the_Bride_Part_II>  7.884109  \n",
       "4                <Heat_(1995_film)>  8.272855  \n",
       "...                             ...       ...  \n",
       "7392           <All_at_Once_(film)>  0.477121  \n",
       "7393      <The_Miracle_(2009_film)>  4.704631  \n",
       "7394          <Savages_(2006_film)>  6.123398  \n",
       "7395      <Pro_Lyuboff_(2010_film)>  6.103391  \n",
       "7396          <Antidur_(2007_film)>  6.150142  \n",
       "\n",
       "[7397 rows x 20 columns]"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "df.to_parquet(Path(yadl_data_path, \"movies_large-yadl.parquet\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_yadl = pl.read_parquet(Path(yadl_data_path, \"movies_large-yadl.parquet\"))\n",
    "\n",
    "df_yadl.select(pl.col(\"target\"), pl.col(\"col_to_embed\")).write_parquet(\n",
    "    Path(yadl_data_path, \"movies_large-yadl-depleted.parquet\")\n",
    ")\n",
    "\n",
    "df_vote = df_yadl.with_columns(target=pl.col(\"vote_average\")).drop(\"vote_average\")\n",
    "df_vote.write_parquet(Path(yadl_data_path, \"movies_vote_large-yadl.parquet\"))\n",
    "df_vote.select(pl.col(\"target\"), pl.col(\"col_to_embed\")).write_parquet(\n",
    "    Path(yadl_data_path, \"movies_vote_large-yadl-depleted.parquet\")\n",
    ")\n",
    "\n",
    "df_yadl.select(pl.col(\"original_title\"), pl.col(\"target\")).write_parquet(\n",
    "    Path(od_data_path, \"movies_large-depleted-open_data.parquet\")\n",
    ")\n",
    "\n",
    "df_vote.select(pl.col(\"original_title\"), pl.col(\"target\")).write_parquet(\n",
    "    Path(od_data_path, \"movies_vote_large-depleted-open_data.parquet\")\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(7397, 19)"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df_vote.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "PosixPath('data/source_tables/yadl')"
      ]
     },
     "execution_count": 32,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "yadl_data_path"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# US Accidents improved"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "df = pl.read_parquet(\"data/source_tables/us_accidents.parquet\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "df = (\n",
    "    df.with_columns(\n",
    "        County=pl.col(\"raw_entities\")\n",
    "        .str.split(\",\")\n",
    "        .list.to_struct()\n",
    "        .struct.rename_fields([\"County\", \"State\"])\n",
    "    )\n",
    "    .unnest(\"County\")\n",
    "    .drop(\"raw_entities\", \"yago4_col_to_embed\")\n",
    "    .rename({\"yago3_col_to_embed\": \"col_to_embed\"})\n",
    "    .with_columns(\n",
    "        col_to_embed=\"<\" + pl.col(\"col_to_embed\") + \">\"\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_2021 = df.filter(pl.col(\"Year\") == 2021).drop(\"Year\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_2021.write_parquet(Path(od_data_path, \"us_accidents_2021-open_data.parquet\"))\n",
    "df_2021.write_parquet(Path(yadl_data_path, \"us_accidents_2021-yadl.parquet\"))\n",
    "df_2021.select(\"County\", \"State\",\"target\").write_parquet(Path(od_data_path, \"us_accidents_2021-depleted-open_data_County.parquet\"))\n",
    "df_2021.select(\"col_to_embed\",\"target\").write_parquet(Path(yadl_data_path, \"us_accidents_2021-yadl-depleted.parquet\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [],
   "source": [
    "df.write_parquet(Path(od_data_path, \"us_accidents_large-open_data.parquet\"))\n",
    "df.write_parquet(Path(yadl_data_path, \"us_accidents_large-yadl.parquet\"))\n",
    "df.select(\"County\", \"State\",\"Year\",\"target\").write_parquet(Path(od_data_path, \"us_accidents_large-depleted-open_data_County.parquet\"))\n",
    "df.select(\"Year\",\"col_to_embed\",\"target\").write_parquet(Path(yadl_data_path, \"us_accidents_large-yadl-depleted.parquet\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Company employees"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "df = pl.read_parquet(Path(raw_data_path, \"company_employees.parquet\"))\n",
    "df.write_parquet(Path(od_data_path, \"company_employees-open_data.parquet\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_yadl = pl.read_parquet(Path(yadl_data_path, \"company_employees-yadl.parquet\"))\n",
    "df_yadl.select(pl.col(\"target\"), pl.col(\"col_to_embed\")).write_parquet(\n",
    "    Path(yadl_data_path, \"company_employees-depleted-yadl.parquet\")\n",
    ")\n",
    "\n",
    "df.select(pl.col(\"name\"), pl.col(\"target\")).write_parquet(\n",
    "    Path(od_data_path, \"company_employees-depleted-open_data.parquet\")\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Housing prices"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "df = pl.read_parquet(Path(yadl_data_path, \"housing_prices-yadl.parquet\"))\n",
    "df.write_parquet(Path(yadl_data_path, \"housing_prices-yadl.parquet\"))\n",
    "df.write_parquet(Path(raw_data_path, \"housing_prices-yadl.parquet\"))\n",
    "df.drop(\"col_to_embed\").write_parquet(\n",
    "    Path(od_data_path, \"housing_prices-open_data.parquet\")\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div><style>\n",
       ".dataframe > thead > tr > th,\n",
       ".dataframe > tbody > tr > td {\n",
       "  text-align: right;\n",
       "  white-space: pre-wrap;\n",
       "}\n",
       "</style>\n",
       "<small>shape: (22_250, 11)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>RegionID</th><th>SizeRank</th><th>City</th><th>RegionType</th><th>StateName</th><th>Code</th><th>Metro</th><th>County</th><th>State</th><th>target</th><th>col_to_embed</th></tr><tr><td>i64</td><td>i64</td><td>str</td><td>str</td><td>str</td><td>str</td><td>str</td><td>str</td><td>str</td><td>f64</td><td>str</td></tr></thead><tbody><tr><td>6181</td><td>0</td><td>&quot;New York&quot;</td><td>&quot;city&quot;</td><td>&quot;NY&quot;</td><td>&quot;NY&quot;</td><td>&quot;New York-Newar…</td><td>&quot;Queens County&quot;</td><td>&quot;New York&quot;</td><td>5.854972</td><td>&quot;&lt;New_York_City…</td></tr><tr><td>17222</td><td>43</td><td>&quot;Buffalo&quot;</td><td>&quot;city&quot;</td><td>&quot;NY&quot;</td><td>&quot;NY&quot;</td><td>&quot;Buffalo-Cheekt…</td><td>&quot;Erie County&quot;</td><td>&quot;New York&quot;</td><td>5.334042</td><td>&quot;&lt;Buffalo,_New_…</td></tr><tr><td>832063</td><td>45</td><td>&quot;Rochester&quot;</td><td>&quot;city&quot;</td><td>&quot;NY&quot;</td><td>&quot;NY&quot;</td><td>&quot;Rochester, NY&quot;</td><td>&quot;Monroe County&quot;</td><td>&quot;New York&quot;</td><td>5.322101</td><td>&quot;&lt;Rochester,_Ne…</td></tr><tr><td>34937</td><td>145</td><td>&quot;Yonkers&quot;</td><td>&quot;city&quot;</td><td>&quot;NY&quot;</td><td>&quot;NY&quot;</td><td>&quot;New York-Newar…</td><td>&quot;Westchester Co…</td><td>&quot;New York&quot;</td><td>5.779882</td><td>&quot;&lt;Yonkers,_New_…</td></tr><tr><td>7353</td><td>149</td><td>&quot;Syracuse&quot;</td><td>&quot;city&quot;</td><td>&quot;NY&quot;</td><td>&quot;NY&quot;</td><td>&quot;Syracuse, NY&quot;</td><td>&quot;Onondaga Count…</td><td>&quot;New York&quot;</td><td>5.233641</td><td>&quot;&lt;Syracuse,_New…</td></tr><tr><td>40779</td><td>229</td><td>&quot;Schenectady&quot;</td><td>&quot;city&quot;</td><td>&quot;NY&quot;</td><td>&quot;NY&quot;</td><td>&quot;Albany-Schenec…</td><td>&quot;Schenectady Co…</td><td>&quot;New York&quot;</td><td>5.435997</td><td>&quot;&lt;Schenectady,_…</td></tr><tr><td>37074</td><td>246</td><td>&quot;Albany&quot;</td><td>&quot;city&quot;</td><td>&quot;NY&quot;</td><td>&quot;NY&quot;</td><td>&quot;Albany-Schenec…</td><td>&quot;Albany County&quot;</td><td>&quot;New York&quot;</td><td>5.455378</td><td>&quot;&lt;Albany,_New_Y…</td></tr><tr><td>34819</td><td>698</td><td>&quot;White Plains&quot;</td><td>&quot;city&quot;</td><td>&quot;NY&quot;</td><td>&quot;NY&quot;</td><td>&quot;New York-Newar…</td><td>&quot;Westchester Co…</td><td>&quot;New York&quot;</td><td>5.825207</td><td>&quot;&lt;White_Plains,…</td></tr><tr><td>831538</td><td>699</td><td>&quot;Binghamton&quot;</td><td>&quot;city&quot;</td><td>&quot;NY&quot;</td><td>&quot;NY&quot;</td><td>&quot;Binghamton, NY…</td><td>&quot;Broome County&quot;</td><td>&quot;New York&quot;</td><td>5.200401</td><td>&quot;&lt;Binghamton,_N…</td></tr><tr><td>26114</td><td>705</td><td>&quot;New Rochelle&quot;</td><td>&quot;city&quot;</td><td>&quot;NY&quot;</td><td>&quot;NY&quot;</td><td>&quot;New York-Newar…</td><td>&quot;Westchester Co…</td><td>&quot;New York&quot;</td><td>5.908062</td><td>&quot;&lt;New_Rochelle,…</td></tr><tr><td>32991</td><td>735</td><td>&quot;Mount Vernon&quot;</td><td>&quot;city&quot;</td><td>&quot;NY&quot;</td><td>&quot;NY&quot;</td><td>&quot;New York-Newar…</td><td>&quot;Westchester Co…</td><td>&quot;New York&quot;</td><td>5.709765</td><td>&quot;&lt;Mount_Vernon,…</td></tr><tr><td>21017</td><td>741</td><td>&quot;Utica&quot;</td><td>&quot;city&quot;</td><td>&quot;NY&quot;</td><td>&quot;NY&quot;</td><td>&quot;Utica-Rome, NY…</td><td>&quot;Oneida County&quot;</td><td>&quot;New York&quot;</td><td>5.23566</td><td>&quot;&lt;Utica,_New_Yo…</td></tr><tr><td>&hellip;</td><td>&hellip;</td><td>&hellip;</td><td>&hellip;</td><td>&hellip;</td><td>&hellip;</td><td>&hellip;</td><td>&hellip;</td><td>&hellip;</td><td>&hellip;</td><td>&hellip;</td></tr><tr><td>34115</td><td>26285</td><td>&quot;Stannard&quot;</td><td>&quot;city&quot;</td><td>&quot;VT&quot;</td><td>&quot;VT&quot;</td><td>null</td><td>&quot;Caledonia Coun…</td><td>&quot;Vermont&quot;</td><td>5.422473</td><td>&quot;&lt;Stannard,_Ver…</td></tr><tr><td>38827</td><td>26341</td><td>&quot;Guildhall&quot;</td><td>&quot;city&quot;</td><td>&quot;VT&quot;</td><td>&quot;VT&quot;</td><td>null</td><td>&quot;Essex County&quot;</td><td>&quot;Vermont&quot;</td><td>5.427118</td><td>&quot;&lt;Guildhall,_Ve…</td></tr><tr><td>44916</td><td>26752</td><td>&quot;Derby Line&quot;</td><td>&quot;city&quot;</td><td>&quot;VT&quot;</td><td>&quot;VT&quot;</td><td>null</td><td>&quot;Orleans County…</td><td>&quot;Vermont&quot;</td><td>5.342342</td><td>&quot;&lt;Derby_Line,_V…</td></tr><tr><td>10246</td><td>26874</td><td>&quot;Baltimore&quot;</td><td>&quot;city&quot;</td><td>&quot;VT&quot;</td><td>&quot;VT&quot;</td><td>&quot;Lebanon, NH-VT…</td><td>&quot;Windsor County…</td><td>&quot;Vermont&quot;</td><td>5.5295</td><td>&quot;&lt;Baltimore,_Ve…</td></tr><tr><td>50400</td><td>27025</td><td>&quot;Woodford&quot;</td><td>&quot;city&quot;</td><td>&quot;VT&quot;</td><td>&quot;VT&quot;</td><td>&quot;Bennington, VT…</td><td>&quot;Bennington Cou…</td><td>&quot;Vermont&quot;</td><td>5.372364</td><td>&quot;&lt;Woodford,_Ver…</td></tr><tr><td>46917</td><td>27122</td><td>&quot;Orleans&quot;</td><td>&quot;city&quot;</td><td>&quot;VT&quot;</td><td>&quot;VT&quot;</td><td>null</td><td>&quot;Orleans County…</td><td>&quot;Vermont&quot;</td><td>5.219014</td><td>&quot;&lt;Orleans,_Verm…</td></tr><tr><td>6240</td><td>27452</td><td>&quot;North Troy&quot;</td><td>&quot;city&quot;</td><td>&quot;VT&quot;</td><td>&quot;VT&quot;</td><td>null</td><td>&quot;Orleans County…</td><td>&quot;Vermont&quot;</td><td>5.173568</td><td>&quot;&lt;North_Troy,_V…</td></tr><tr><td>396478</td><td>27594</td><td>&quot;Landgrove&quot;</td><td>&quot;city&quot;</td><td>&quot;VT&quot;</td><td>&quot;VT&quot;</td><td>&quot;Bennington, VT…</td><td>&quot;Bennington Cou…</td><td>&quot;Vermont&quot;</td><td>5.912688</td><td>&quot;&lt;Landgrove,_Ve…</td></tr><tr><td>27831</td><td>27951</td><td>&quot;Wells River&quot;</td><td>&quot;city&quot;</td><td>&quot;VT&quot;</td><td>&quot;VT&quot;</td><td>&quot;Lebanon, NH-VT…</td><td>&quot;Orange County&quot;</td><td>&quot;Vermont&quot;</td><td>5.369254</td><td>&quot;&lt;Wells_River,_…</td></tr><tr><td>398478</td><td>28081</td><td>&quot;Lemington&quot;</td><td>&quot;city&quot;</td><td>&quot;VT&quot;</td><td>&quot;VT&quot;</td><td>null</td><td>&quot;Essex County&quot;</td><td>&quot;Vermont&quot;</td><td>5.349301</td><td>&quot;&lt;Lemington,_Ve…</td></tr><tr><td>56021</td><td>28699</td><td>&quot;Jeffersonville…</td><td>&quot;city&quot;</td><td>&quot;VT&quot;</td><td>&quot;VT&quot;</td><td>null</td><td>&quot;Lamoille Count…</td><td>&quot;Vermont&quot;</td><td>5.559919</td><td>&quot;&lt;Jeffersonvill…</td></tr><tr><td>249186</td><td>28699</td><td>&quot;Derby Center&quot;</td><td>&quot;city&quot;</td><td>&quot;VT&quot;</td><td>&quot;VT&quot;</td><td>null</td><td>&quot;Orleans County…</td><td>&quot;Vermont&quot;</td><td>5.396455</td><td>&quot;&lt;Derby_Center,…</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (22_250, 11)\n",
       "┌──────────┬──────────┬────────────┬────────────┬───┬────────────┬──────────┬──────────┬───────────┐\n",
       "│ RegionID ┆ SizeRank ┆ City       ┆ RegionType ┆ … ┆ County     ┆ State    ┆ target   ┆ col_to_em │\n",
       "│ ---      ┆ ---      ┆ ---        ┆ ---        ┆   ┆ ---        ┆ ---      ┆ ---      ┆ bed       │\n",
       "│ i64      ┆ i64      ┆ str        ┆ str        ┆   ┆ str        ┆ str      ┆ f64      ┆ ---       │\n",
       "│          ┆          ┆            ┆            ┆   ┆            ┆          ┆          ┆ str       │\n",
       "╞══════════╪══════════╪════════════╪════════════╪═══╪════════════╪══════════╪══════════╪═══════════╡\n",
       "│ 6181     ┆ 0        ┆ New York   ┆ city       ┆ … ┆ Queens     ┆ New York ┆ 5.854972 ┆ <New_York │\n",
       "│          ┆          ┆            ┆            ┆   ┆ County     ┆          ┆          ┆ _City>    │\n",
       "│ 17222    ┆ 43       ┆ Buffalo    ┆ city       ┆ … ┆ Erie       ┆ New York ┆ 5.334042 ┆ <Buffalo, │\n",
       "│          ┆          ┆            ┆            ┆   ┆ County     ┆          ┆          ┆ _New_York │\n",
       "│          ┆          ┆            ┆            ┆   ┆            ┆          ┆          ┆ >         │\n",
       "│ 832063   ┆ 45       ┆ Rochester  ┆ city       ┆ … ┆ Monroe     ┆ New York ┆ 5.322101 ┆ <Rocheste │\n",
       "│          ┆          ┆            ┆            ┆   ┆ County     ┆          ┆          ┆ r,_New_Yo │\n",
       "│          ┆          ┆            ┆            ┆   ┆            ┆          ┆          ┆ rk>       │\n",
       "│ 34937    ┆ 145      ┆ Yonkers    ┆ city       ┆ … ┆ Westcheste ┆ New York ┆ 5.779882 ┆ <Yonkers, │\n",
       "│          ┆          ┆            ┆            ┆   ┆ r County   ┆          ┆          ┆ _New_York │\n",
       "│          ┆          ┆            ┆            ┆   ┆            ┆          ┆          ┆ >         │\n",
       "│ …        ┆ …        ┆ …          ┆ …          ┆ … ┆ …          ┆ …        ┆ …        ┆ …         │\n",
       "│ 27831    ┆ 27951    ┆ Wells      ┆ city       ┆ … ┆ Orange     ┆ Vermont  ┆ 5.369254 ┆ <Wells_Ri │\n",
       "│          ┆          ┆ River      ┆            ┆   ┆ County     ┆          ┆          ┆ ver,_Verm │\n",
       "│          ┆          ┆            ┆            ┆   ┆            ┆          ┆          ┆ ont>      │\n",
       "│ 398478   ┆ 28081    ┆ Lemington  ┆ city       ┆ … ┆ Essex      ┆ Vermont  ┆ 5.349301 ┆ <Lemingto │\n",
       "│          ┆          ┆            ┆            ┆   ┆ County     ┆          ┆          ┆ n,_Vermon │\n",
       "│          ┆          ┆            ┆            ┆   ┆            ┆          ┆          ┆ t>        │\n",
       "│ 56021    ┆ 28699    ┆ Jeffersonv ┆ city       ┆ … ┆ Lamoille   ┆ Vermont  ┆ 5.559919 ┆ <Jefferso │\n",
       "│          ┆          ┆ ille       ┆            ┆   ┆ County     ┆          ┆          ┆ nville,_V │\n",
       "│          ┆          ┆            ┆            ┆   ┆            ┆          ┆          ┆ ermont>   │\n",
       "│ 249186   ┆ 28699    ┆ Derby      ┆ city       ┆ … ┆ Orleans    ┆ Vermont  ┆ 5.396455 ┆ <Derby_Ce │\n",
       "│          ┆          ┆ Center     ┆            ┆   ┆ County     ┆          ┆          ┆ nter,_Ver │\n",
       "│          ┆          ┆            ┆            ┆   ┆            ┆          ┆          ┆ mont>     │\n",
       "└──────────┴──────────┴────────────┴────────────┴───┴────────────┴──────────┴──────────┴───────────┘"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# US Elections"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "df = pl.read_parquet(Path(yadl_data_path, \"us_elections-yadl.parquet\"))\n",
    "df_dem = df.filter(pl.col(\"party\") == \"DEMOCRAT\")\n",
    "df_dem.write_parquet(Path(yadl_data_path, \"us_elections_dem-yadl.parquet\"))\n",
    "df_dem.drop(\"col_to_embed\").write_parquet(\n",
    "    Path(od_data_path, \"us_elections_dem-open_data.parquet\")\n",
    ")\n",
    "df_dem.select(pl.col(\"target\"), pl.col(\"col_to_embed\")).write_parquet(\n",
    "    Path(yadl_data_path, \"us_elections_dem-yadl-depleted.parquet\")\n",
    ")\n",
    "df_dem.select(pl.col(\"target\"), pl.col(\"county_name\")).write_parquet(\n",
    "    Path(od_data_path, \"us_elections_dem-open_data-depleted.parquet\")\n",
    ")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "bench-repro",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
